import requests
from requests.models import PreparedRequest
from PIL import Image
import numpy as np
import torch
from torchvision.transforms import ToPILImage
from io import BytesIO
import os
import time
from datetime import datetime
import base64
import io

API_KEY = os.environ.get("SAI_API_KEY")

# Check for API key in file as a backup, not recommended
try:
    if not API_KEY:
        dir_path = os.path.dirname(os.path.realpath(__file__))
        with open(os.path.join(dir_path, "sai_platform_key.txt"), "r") as f:
            API_KEY = f.read().strip()
            print(f"API Key found in sai_platform_key.txt: {API_KEY}")
        # Validate the key is not empty
        if API_KEY.strip() == "":
            raise Exception(f"API Key is required to use the Stability API. \nPlease set the SAI_API_KEY environment variable to your API key or place in {dir_path}/sai_platform_key.txt.")
        
except Exception as e:
    print(f"\n\n***API Key is required to use the Stability API. Please set the SAI_API_KEY environment variable to your API key or place in {dir_path}/sai_platform_key.txt.***\n\n")

ROOT_API = "https://api.stability.ai/v2beta/"


class StabilityBase:
    API_ENDPOINT = ""
    POLL_ENDPOINT = ""
    ACCEPT = ""

    @classmethod
    def INPUT_TYPES(cls):
        return cls.INPUT_SPEC

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "call"
    CATEGORY = "AI WizArt/Stability AI Suite"

    def call(self, *args, **kwargs):
        
        buffered = BytesIO()
        files = {'none': None}
        data = None
        
        image = kwargs.get('image', None)
        if image is not None:
            kwargs["mode"] = "image-to-image"
            kwargs.pop("aspect_ratio", None)
            image = ToPILImage()(image.squeeze(0).permute(2,0,1))
            image.save(buffered, format="PNG")
            files = self._get_files(buffered, **kwargs)
        else:
            kwargs.pop("strength", None)
        
        style = kwargs.get('style', False)
        if style is False:
            kwargs.pop('style_preset', None)

        headers = {
            "Authorization": API_KEY,
        }

        if kwargs.get("api_key_override"):
            headers = {
                "Authorization": kwargs.get("api_key_override"),
            }

        if headers.get("Authorization") is None:
            raise Exception(f"No Stability key set.\n\nUse your Stability AI API key by:\n1. Setting the SAI_API_KEY environment variable to your API key\n3. Placing inside sai_platform_key.txt\n4. Passing the API key as an argument to the function with the key 'api_key_override'")

        headers["Accept"] = self.ACCEPT

        data = self._get_data(**kwargs)

        req = PreparedRequest()
        req.prepare_method('POST')
        req.prepare_url(f"{ROOT_API}{self.API_ENDPOINT}", None)
        req.prepare_headers(headers)
        req.prepare_body(data=data, files=files)
        response = requests.Session().send(req)

        if response.status_code == 200:
            if self.POLL_ENDPOINT != "":
                id = response.json().get("id")
                logFile(f"Image/video ID for recovery: {id}") # saving id for recovery in case of malfunction
                timeout = 240
                start_time = time.time()
                while True:
                    response = requests.get(f"{ROOT_API}{self.POLL_ENDPOINT}{id}", headers=headers)
                    if response.status_code == 200:
                        if self.ACCEPT == "image/*":
                            return self._return_image(response)
                        if self.ACCEPT == "video/*":
                            return self._return_video(response)
                        break
                    elif response.status_code == 202:
                        time.sleep(10)
                    elif time.time() - start_time > timeout:
                        raise Exception("Stability API Timeout: Request took too long to complete")
                    else:
                        error_info = response.json()
                        raise Exception(f"Stability API Error: {error_info}")
            else:
                result_image = Image.open(BytesIO(response.content))
                result_image = result_image.convert("RGBA")
                result_image = np.array(result_image).astype(np.float32) / 255.0
                result_image = torch.from_numpy(result_image)[None,]
                return (result_image,)
        else:
            error_info = response.json()
            if error_info.get("name") == "unauthorized":
                raise Exception("Stability API Error: Unauthorized.\n\nUse your Stability AI API key by:\n1. Setting the SAI_API_KEY environment variable to your API key\n3. Placing inside sai_platform_key.txt\n4. Passing the API key as an argument to the function with the key 'api_key_override'")
            if error_info.get("name") == "payment_required":
                raise Exception("Stability API Error: Not enough credits.\n\nPlease ensure your SAI API account has enough credits to complete this action.")
            if error_info.get("name") == "bad_request":
                errors = '\n'.join(error_info.get('errors'))
                raise Exception(f"Stability API Error: Bad request.\n\n{errors}")
            else:
                raise Exception(f"Stability API Error: {error_info}")
    
    def _return_image(self, response):
        result_image = Image.open(BytesIO(response.content))
        result_image = result_image.convert("RGBA")
        result_image = np.array(result_image).astype(np.float32) / 255.0
        result_image = torch.from_numpy(result_image)[None,]
        return (result_image,)

    def _return_video(self, response):
        result_video = response.content
        return (result_video,)

    def _get_files(self, buffered, **kwargs):
        return {
            "image": buffered.getvalue()
        }

    def _get_data(self, **kwargs):
        return {k: v for k, v in kwargs.items() if k != "image"}


class StabilityCore(StabilityBase):
    API_ENDPOINT = "stable-image/generate/core"
    ACCEPT = "image/*"
    INPUT_SPEC = {
        "required": {
            "prompt": ("STRING", {"multiline": True}),
        },
        "optional": {
            "negative_prompt": ("STRING", {"multiline": True}),
            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
            "output_format": (["png", "webp", "jpeg"],),
            "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],),
            "style": ("BOOLEAN", {"default": False}),
            "style_preset": (["3d-model", "analog-film", "anime", "cinematic", "comic-book", "digital-art", "enhance", "fantasy-art", "isometric", "line-art", "low-poly", "modeling-compound", "neon-punk", "origami", "photographic", "pixel-art", "tile-texture"],),
            "api_key_override": ("STRING", {"multiline": False}),
        }
    }


class StabilityCreativeUpscale(StabilityBase):
    API_ENDPOINT = "stable-image/upscale/creative"
    POLL_ENDPOINT  = "stable-image/upscale/creative/result/"
    ACCEPT = "image/*"
    INPUT_SPEC = {
        "required": {
            "image": ("IMAGE",),
            "prompt": ("STRING", {"multiline": True}),
        },
        "optional": {
            "negative_prompt": ("STRING", {"multiline": True}),
            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
            "creativity": ("FLOAT", {"default": 0.3, "min": 0.01, "max": 0.35, "step": 0.01}),
            "output_format": (["png", "webp", "jpeg"],),
            "api_key_override": ("STRING", {"multiline": False}),
        }
    }


class StabilityRemoveBackground(StabilityBase):
    API_ENDPOINT = "stable-image/edit/remove-background"
    ACCEPT = "image/*"
    INPUT_SPEC = {
        "required": {
            "image": ("IMAGE",),
        },
    }

class StabilityInpainting(StabilityBase):
    API_ENDPOINT = "stable-image/edit/inpaint"
    ACCEPT = "image/*"
    INPUT_SPEC = {
        "required": {
            "image": ("IMAGE",),
            "mask": ("MASK",),
            "prompt": ("STRING", {"multiline": True, "default": ""}),\
        },
        "optional": {
            "negative_prompt": ("STRING", {"multiline": True, "default": ""}),\
            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
            "output_format": (["png", "webp", "jpeg"],),
            "api_key_override": ("STRING", {"multiline": False}),
        }
    }
    def _get_files(self, buffered, **kwargs):
        mask = kwargs.get("mask")
        to_pil = ToPILImage()
        mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
        mask = to_pil(mask.squeeze(0).permute(2,0,1))
        buffered_mask = BytesIO()
        mask.save(buffered_mask, format="PNG")
        return {
            "image": buffered.getvalue(),
            "mask": buffered_mask.getvalue(),
        }


class StabilitySearchAndReplace(StabilityBase):
    API_ENDPOINT = "stable-image/edit/search-and-replace"
    ACCEPT = "image/*"
    INPUT_SPEC = {
        "required": {
            "image": ("IMAGE",),
            "search_prompt": ("STRING", {"multiline": True}, "Search Prompt"),
            "prompt": ("STRING", {"multiline": True}),
        },
        "optional": {
            "negative_prompt": ("STRING", {"multiline": True}),
            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
            "api_key_override": ("STRING", {"multiline": False}),
            "output_format": (["png", "webp", "jpeg"],),
        },
    }


class StabilitySD3(StabilityBase):
    API_ENDPOINT = "stable-image/generate/sd3"
    ACCEPT = "image/*"
    INPUT_SPEC = {
        "required": {
            "model": (["sd3", "sd3-turbo"],),
            "prompt": ("STRING", {"multiline": True}),
        },
        "optional": {
            "image": ("IMAGE",),
            "negative_prompt": ("STRING", {"multiline": True}),
            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
            "strength": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 1.0, "step": 0.01}),
            "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],),
            "output_format": (["png", "jpeg"],),
            "api_key_override": ("STRING", {"multiline": False}),
        },
    }


class StabilityOutpainting(StabilityBase):
    API_ENDPOINT = "stable-image/edit/outpaint"
    ACCEPT = "image/*"
    INPUT_SPEC = {
        "required": {
            "image": ("IMAGE",),
            "left": ("INT", {"default": 0, "min": 0, "max": 512}),
            "right": ("INT", {"default": 0, "min": 0, "max": 512}),
            "up": ("INT", {"default": 0, "min": 0, "max": 512}),
            "down": ("INT", {"default": 0, "min": 0, "max": 512}),
        },
        "optional": {
            "prompt": ("STRING", {"multiline": True}),
            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
            "output_format": (["png", "webp", "jpeg"],),
            "api_key_override": ("STRING", {"multiline": False}),
        },
    }

# ========================================================
# FILE RECOVER
# ========================================================

class StabilityCreativeUpscaleRecover(StabilityBase):
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image_id": ("STRING", {
                    "multiline": False
                })
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image_out",)
    FUNCTION = "creativeUpscaleRecover"

    def creativeUpscaleRecover(self, image_id):
        get_code = 202
        while get_code == 202:
            response_get = requests.request(
                "GET",
                f"https://api.stability.ai/v2beta/stable-image/upscale/creative/result/{image_id}",
                headers={
                    "accept": "application/json",
                    "authorization": f"Bearer {API_KEY}"
                },
            )
            get_code = response_get.status_code
            time.sleep(10)
            print("Waiting image...")
        if response_get.status_code == 200:
            json_data = response_get.json()
            image_base64 = json_data['image']
            image_bytes = base64.b64decode(image_base64)
            image_data = Image.open(io.BytesIO(image_bytes))
            output_t = pil2tensor(image_data)
            return (output_t,)
        else:
            print(response_get.json())

# ========================================================
# UTILITIES
# ========================================================

def logFile(text):
    now = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    new_entry = f"{now} - {text}"
    logfile = os.path.join(dir_path, 'log.txt')
    with open(logfile, "a") as file:
        file.write(new_entry + "\n")

def pil2tensor(image):
    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)